from collections import defaultdict
import itertools

class Apriori:
    def __init__(self, min_support=0.4, min_confidence=0.5):
        self.min_support = min_support  # 最小支持度
        self.min_confidence = min_confidence  # 最小置信度
        self.transactions = []  # 交易数据集
        self.itemsets = {}  # 存储不同大小的频繁项集
        self.rules = []  # 存储关联规则

    def fit(self, transactions):
        self.transactions = transactions
        # 第1步：生成频繁1项集
        itemsets_1 = self._generate_itemsets(1)
        self.itemsets[1] = itemsets_1
        print("频繁1项集:", itemsets_1)  # 打印频繁1项集
        k = 2
        while True:
            # 生成频繁k项集
            itemsets_k = self._generate_itemsets(k)
            if not itemsets_k:
                break
            self.itemsets[k] = itemsets_k
            print(f"频繁{k}项集:", itemsets_k)  # 打印频繁k项集
            k += 1

    def _generate_itemsets(self, k):
        itemsets = defaultdict(int)
        for transaction in self.transactions:
            for itemset in itertools.combinations(transaction, k):
                itemsets[itemset] += 1
        total_transactions = len(self.transactions)
        frequent_itemsets = {itemset for itemset, count in itemsets.items() if count / total_transactions >= self.min_support}
        return frequent_itemsets

    def generate_rules(self):
        for k in range(2, len(self.itemsets) + 1):
            for itemset in self.itemsets[k]:
                for subset in itertools.combinations(itemset, len(itemset) - 1):
                    antecedent = set(subset)
                    consequent = set(itemset) - antecedent
                    antecedent_support = self._support(antecedent)
                    rule_support = self._support(itemset)
                    confidence = rule_support / antecedent_support
                    print(f"Rule: {antecedent} -> {consequent}, Support: {rule_support:.2f}, Confidence: {confidence:.2f}")  # 打印规则的支持度和置信度
                    if confidence >= self.min_confidence:
                        self.rules.append((antecedent, consequent, confidence))

    def _support(self, itemset):
        itemset = set(itemset)
        count = sum(1 for transaction in self.transactions if itemset.issubset(transaction))
        return count / len(self.transactions)

    def get_rules(self):
        return self.rules


# 示例数据集
transactions = [
    {'牛奶', '面包'},
    {'牛奶', '黄油'},
    {'面包', '黄油'},
    {'牛奶', '面包', '黄油'},
    {'面包'}
]

# 创建 Apriori 实例
apriori = Apriori(min_support=0.4, min_confidence=0.5)

# 训练模型
apriori.fit(transactions)

# 生成关联规则
apriori.generate_rules()

# 获取并打印规则
rules = apriori.get_rules()
if not rules:
    print("没有找到符合条件的规则。")
else:
    for rule in rules:
        print(f"Antecedent: {rule[0]}, Consequent: {rule[1]}, Confidence: {rule[2]:.2f}")